% Copyright (C) 2009,2010,2011,2012  Marco Restelli
%
% This file is part of:
%   LDGH -- Local Hybridizable Discontinuous Galerkin toolkit
%
% LDGH is free software: you can redistribute it and/or modify it
% under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% LDGH is distributed in the hope that it will be useful, but WITHOUT
% ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
% or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public
% License for more details.
%
% You should have received a copy of the GNU General Public License
% along with LDGH. If not, see <http://www.gnu.org/licenses/>.
%
% author: Marco Restelli                   <marco.restelli@gmail.com>


function [grids,p_fixed] = grid_partition(p,e,t,idx,varargin)
%
% grids = grid_partition(p,e,t,idx)
% grids = grid_partition(p,e,t,idx,periodic_faces)
% grids = grid_partition(p,e,t,idx,periodic_faces,data_names,data)
% [grids,p_fixed] = grid_partition(...)
%
% Partition the mesh (p,e,t) into length(idx) subgrids using the
% indexes in idx.
%
% Notes:
%
% A) The index vector idx starts from 0: this is consistent with both
% METIS and MPI.
%
% B) The general procedure for partitioning a mesh is:
% 1) write the complete mesh in a form that is readable by mpmetis,
%   that is
%     Nel
%     e1v1 e1v2 e1v3 [e1v4]
%     e2v1 e2v2 e2v3 [e2v4]
%     ...
%   where Nel is the number of elements and the following rows are the
%   nodes of the elements (vertex numbering starts from 1)
% 2) partition the mesh with  "mpmetis -gtype=dual input-file N"
% 3) use this function with the idx provided by mpmetis
% 4) save each mesh returned in grids as a separate mesh file.
%
% C) One could also use "-gtype=nodal", however, since we distribute
% the elements and duplicate nodes and sides, the communications are
% minimized using "-gtpe=dual".
%
% D) The decomposed grid files can be written as follows:
% fname = 'xyz';
% d = 23;
% for i=1:size(grids,2)
%   p=grids(i).p; t=grids(i).t; e=grids(i).e; ivl2ivn=grids(i).ivl2ivn;
%   save([fname,'.',num2str(i-1,'%.3d')],'d','p','e','t','ivl2ivn');
% end
%
% E) The return value grids includes a field ivl2ivn, which is an
% array with one element for each subgrid node. Each element is a
% struct with the following fields (assuming iv is the index of the
% node):
%  ivl2ivn(iv).nd -> number of neighboring grids sharing this vertex
%  ivl2ivn(iv).id(:) -> indexes of the neighboring grids
%  ivl2ivn(iv).iv(:) -> local indexes in the neighboring grids
%
% F) The optional argument periodic_faces is a list of pairs of faces
% wich have periodic boundary conditions. The list is in the form
%    periodic_faces = [ mf1 mf2 ... ]
%                     [ sf1 sf2 ... ]
%                     [ t11 t12 ... ]
%                     [ t21 t22 ... ]
%                     [ ... ... ... ]
% where each column corresponds to a master/slave pair and mf and sf
% are the indexes of two faces (master and slave) and [t1 t2 ...] is
% the translation vector which maps sf on mf.
%
% G) The output argument p_fixed is useful when using periodic
% boundary conditions on a mesh which is not exactly periodic: this is
% in fact a copy of the input value of p where the nodes on the slave
% boundaries are moved so that they correspond to the position of the
% master ones. If one wishes to use this corrected nodes, one should
% call the function a second time using p_fixed instead of p.
%
% H) The optional arguments data_names and data provide nodal data
% which are then partitioned together with the grid and stored in
% additional fields of the return value grids. data_fields must be a
% list of strings with the data names (used to name the corresponding
% additional fields in grids), while data must be an array where each
% column corresponds to a data field. To pass nodal data without
% specifying periodicity information, use
%   grids = grid_partition(p,e,t,idx,[],data_names,data)

 %---------------------------------------------------------------------
 % Check the input arguments

 % periodicity
 if((nargin>4)&&(~isempty(varargin{1})))
   sideper = varargin{1};
   for i=1:size(sideper,2)
     VM{i} = [];
     VS{i} = [];
   end
   toll = 1e-10;
 else
   sideper = zeros(2,0);
 end
 % nodal data
 have_noda_data = 0;
 if(nargin>5)
   have_noda_data = 1;
   node_data_names = varargin{2};
   node_data       = varargin{3};
   if(numel(node_data_names)~=columns(node_data))
     error(['The number of data names is not consistent ' ...
            'with the number of data fields']);
   end
 end
 %---------------------------------------------------------------------

 % initialize working arrays
 N = max(idx)+1;  % number of subgrids (idx starts from 0)
 dim = size(p,1); % space dimension
 iv_ren = zeros(N,size(p,2)); % local renumbering for each subgrid
 nv = zeros(N,1); % number of nodes for each subgrid
 nt = zeros(N,1); % number of elements for each subgrid

 % loop over the elements to count the local nodes and elements
 for ie=1:size(t,2)
   isg = idx(ie)+1; % index of the subgrid (starting from 1)
   nt(isg) = nt(isg) + 1;
   for l=1:dim+1 % loop on the nodes
     iv = t(l,ie);
     if(iv_ren(isg,iv)==0) % new local vertex
       nv(isg) = nv(isg) + 1;
       iv_ren(isg,iv) = nv(isg);
     end
   end
 end

 % allocate the local grids
 for isg=1:N
   grids(isg) = struct( ...
     'p', zeros(dim,nv(isg)), ...
     'e', zeros(size(e)), ... % will be reshaped
     't', zeros(dim+1,nt(isg)), ...
     'ivl2ivn', [] );
   grids(isg).ivl2ivn(1:nv(isg)) = struct( ...
     'nd', 0, ...
     'id', [], ...
     'iv', [] );
 end

 % fill grids.p
 for iv=1:size(p,2)
   for isg=1:N
     if(iv_ren(isg,iv)!=0) % grid isg sees vertex iv
       grids(isg).p(:,iv_ren(isg,iv)) = p(:,iv);
     end
   end
 end

 % fill grids.t
 nt(:) = 0; % reset the counter
 for ie=1:size(t,2)
   isg = idx(ie)+1; % index of the subgrid
   nt(isg) = nt(isg) + 1;

   % insert the nodes
   iv = t(1:dim+1,ie);                       % global indexes
   grids(isg).t(:,nt(isg)) = iv_ren(isg,iv); % local indexes
 end

 % fill grids.ivl2ivn (periodicity information will be added later)
 for iv=1:size(p,2)
   nsg = sum( (iv_ren(:,iv)!=0) ); % iv belongs to nsg subgrids
   if(nsg>1) % the node belongs to more than one grid
     isg = find( (iv_ren(:,iv)!=0) );
     for i=1:length(isg)
       ign = setdiff(isg,isg(i)); % neighboring subgrids
       grids(isg(i)).ivl2ivn(iv_ren(isg(i),iv)) = struct( ...
         'nd', nsg-1, ...
         'id', ign-1, ... % subdomains are numbered starting from 0
         'iv', iv_ren(ign,iv)' ); % local indexes on neigh. subgrids
     end
   end
 end

 % partition e and include periodicity
 % This is done in the following steps:
 %  1) identify the boundary nodes
 %  2) define a vertex to elements pointer for the boundary nodes
 %  3) assign the boundary sides to the subgrids
 v2bv = zeros(size(p,2),1); % vertex to boundary vertex
 nbv  = 0;                  % number of boundary vertexes
 for ibs=1:size(e,2)
   % Check whether the face to which the boundary side belongs is
   % included in the list of master/slave periodic faces. We assume
   % that a face can appear in sideper only once.
   [i_ms,i_sp] = find(sideper(1:2,:)==e(dim+2+1,ibs));
   if(isempty(i_ms)) % side is not periodic
     for ivl=1:dim
       iv = e(ivl,ibs);
       if(v2bv(iv)==0) % found a new boundary vertex
         nbv = nbv + 1;
         v2bv(iv) = nbv;
       end
     end
   else % side is periodic: add the connectivity information
     if(i_ms==1) % side belongs to the master face
       VM{i_sp} = [VM{i_sp},e(1:dim,ibs)']; % list of nodes (master)
     else % slave
       VS{i_sp} = [VS{i_sp},e(1:dim,ibs)']; % list of nodes (slave)
     end
     % mark the side, so that it will be no longer considered a
     % boundary side
     e(1,ibs) = -1;
   end
 end
 if(nbv>0)
   bv2e{nbv} = []; % boundary vertex -> connected elements
 end
 for ie=1:size(t,2)
   for ivl=1:dim+1
     ibv = v2bv(t(ivl,ie));
     if(ibv!=0)
       bv2e{ibv} = [bv2e{ibv} , ie];
     end
   end
 end
 nbs = zeros(N,1);
 for ibs=1:size(e,2)
   if(e(1,ibs)!=-1) % side is not marked
     % identify the connected element
     ie = bv2e{v2bv(e(1,ibs))};
     for ivl=2:dim
       ie = intersect(ie,bv2e{v2bv(e(ivl,ibs))});
     end
     if(length(ie)>1) % we should never get here
       error(['The boundary side ',num2str(ibs), ...
              ' belongs to more than one element: ',num2str(ie)]);
     end
     % now we know that the boundary side belongs to ie
     bs = e(:,ibs);
     isg = idx(ie) + 1; % subgrid index
     bs(1:dim,1) = iv_ren(isg,bs(1:dim,1)); % local node numbering
     nbs(isg) = nbs(isg) + 1;
     grids(isg).e(:,nbs(isg)) = bs; % assign the boundary side to isg
   end
 end
 for isg=1:N
   grids(isg).e = grids(isg).e(:,1:nbs(isg)); % trim grids.e
 end

 % fix the periodicity information
 p_fixed = p;
 for i=1:size(sideper,2)  % loop on the master faces
   vm = unique(VM{i}); % master nodes for this face (global numbering)
   vs = unique(VS{i}); % slave nodes for this face (global numbering)
   found_slave = zeros(size(vs));
   for ivm=1:length(vm)   % loop on the master nodes
     xvm = p(:,vm(ivm)) - sideper(3:end,i); % periodic translation
     found_master = 0; dist = inf;
     for ivs=1:length(vs) % loop on the slave nodes
       distn = norm(p(:,vs(ivs))-xvm);
       if(distn<=toll) % found
         found_master = 1;
         found_slave(ivs) = 1;
         
         % domains on the master side
         dom_ms = find(iv_ren(:,vm(ivm))!=0); % subdomain(s)
         num_ms = iv_ren(dom_ms,vm(ivm));     % local indexes

         % domains on the slave side
         dom_ss = find(iv_ren(:,vs(ivs))!=0); % subdomain(s)
         num_ss = iv_ren(dom_ss,vs(ivs));     % local indexes

         % add data
         for iv=1:length(dom_ms)
           grids(dom_ms(iv)).ivl2ivn(num_ms(iv)).nd = ...
             grids(dom_ms(iv)).ivl2ivn(num_ms(iv)).nd + length(dom_ss);
           grids(dom_ms(iv)).ivl2ivn(num_ms(iv)).id = ...
             [grids(dom_ms(iv)).ivl2ivn(num_ms(iv)).id , dom_ss'-1];
           grids(dom_ms(iv)).ivl2ivn(num_ms(iv)).iv = ...
             [grids(dom_ms(iv)).ivl2ivn(num_ms(iv)).iv , num_ss'];
         end
         for iv=1:length(dom_ss)
           grids(dom_ss(iv)).ivl2ivn(num_ss(iv)).nd = ...
             grids(dom_ss(iv)).ivl2ivn(num_ss(iv)).nd + length(dom_ms);
           grids(dom_ss(iv)).ivl2ivn(num_ss(iv)).id = ...
             [grids(dom_ss(iv)).ivl2ivn(num_ss(iv)).id , dom_ms'-1];
           grids(dom_ss(iv)).ivl2ivn(num_ss(iv)).iv = ...
             [grids(dom_ss(iv)).ivl2ivn(num_ss(iv)).iv , num_ms'];
         end

         break
       else
         if(distn<dist)
           dist = distn;
           i_nearest = vs(ivs);
         end
       end
     end
     if(~found_master)
       warning([ ...
         'Unable to find a corresponding node for iv=', ...
                                              num2str(vm(ivm)), ...
         ', belonging to boundary face ',num2str(sideper(1,i)), ...
         ', periodic with boundary face ',num2str(sideper(2,i)),...
         '; nearest node: ',num2str(i_nearest), ...
         ', distance: ',num2str(dist)]);
       p_fixed(:,i_nearest) = p(:,vm(ivm)) - sideper(3:end,i);
     end
   end
   if(~isempty(find(found_slave==0)))
     warning([ ...
     'Unable to find a corresponding node for the following nodes: ',...
     num2str(vs(find(found_slave==0))), ...
     ', belonging to boundary face ',num2str(sideper(2,i)), ...
     ', periodic with boundary face ',num2str(sideper(1,i)) ]);
   end
 end

 % partition nodal data
 if(have_noda_data)
   for isg=1:N % loop on subgrids
     for id=1:numel(node_data_names) % loop on data fields
       [~,globa_idx,local_idx] = find(iv_ren(isg,:));
       grids(isg).(node_data_names{id})(local_idx) = ...
                              node_data(globa_idx,id);
     end
   end
 end

return

